﻿/*
 * @author: wizardc
 */

using System.Collections.Generic;
using Dou.Utils;
using Unity.Collections;
using Unity.Collections.LowLevel.Unsafe;
using Unity.Mathematics;
using UnityEngine;
using UnityEngine.Rendering;

namespace Dou.GPU.Instancing
{
    /// <summary>
    /// GPUInstancing 批次绘制类
    /// </summary>
    public unsafe partial class GPUBatch
    {
        private const float Range = 10000f;
        private static readonly Bounds BOUNDS = new Bounds(Vector3.zero, new Vector3(Range, Range, Range));
        
        private Mesh _mesh;
        private uint _meshIndexCount;
        private uint _meshIndexStart;
        private uint _meshBaseVertex;
        private Material _material;
        private int _maximum;
        private ShadowCastingMode _castShadows;
        private bool _receiveShadows;
        private int _layer;
        private Camera _camera;
        
        private List<GPUPropertyData> _propertyDataList = new List<GPUPropertyData>();
        private List<GPUBatchData> _batchDataList = new List<GPUBatchData>();

        public void Init(Mesh mesh, Material material, int maximum, ShadowCastingMode castShadows, bool receiveShadows, int layer, Camera camera)
        {
            _mesh = mesh;
            _meshIndexCount = mesh.GetIndexCount(0);
            _meshIndexStart = mesh.GetIndexStart(0);
            _meshBaseVertex = mesh.GetBaseVertex(0);
            _material = material;
            _maximum = maximum;
            _castShadows = castShadows;
            _receiveShadows = receiveShadows;
            _layer = layer;
            _camera = camera;
        }

        private GPUPropertyData GetPropertyData(int nameId)
        {
            for (var i = 0; i < _propertyDataList.Count; i++)
            {
                if (_propertyDataList[i].nameId == nameId)
                {
                    return _propertyDataList[i];
                }
            }
            return null;
        }

        public void RegisterProperty(int nameId, int stride)
        {
            var data = GetPropertyData(nameId);
            if (data == null)
            {
                data = new GPUPropertyData();
                data.nameId = nameId;
                data.stride = stride;
                data.nativeList = new NativeList<byte>(Allocator.Persistent);
                _propertyDataList.Add(data);
            }
        }

        public GPUPropertyData GetProperty(int nameId)
        {
            var data = GetPropertyData(nameId);
            return data;
        }
        
        #region 提供常见数据类型设置方法

        #region Int
        
        public void SetInt(int nameId, int[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (int* ptr = &data[0])
                {
                    property.nativeList.Clear();
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Int);
                }
            }
        }

        public void SetInt(int nameId, List<int> data)
        {
            SetInt(nameId, data.ToArray());
        }
        
        public void SetInt(int nameId, ref NativeArray<int> data)
        {
            SetInt(nameId, (int*) data.GetUnsafePtr(), data.Length);
        }

        public void SetInt(int nameId, ref NativeList<int> data)
        {
            SetInt(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void SetInt(int nameId, int* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.Clear();
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Int);
            }
        }

        public void AppendInt(int nameId, int[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (int* ptr = &data[0])
                {
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Int);
                }
            }
        }

        public void AppendInt(int nameId, List<int> data)
        {
            AppendInt(nameId, data.ToArray());
        }

        public void AppendInt(int nameId, ref NativeArray<int> data)
        {
            AppendInt(nameId, (int*) data.GetUnsafePtr(), data.Length);
        }
        
        public void AppendInt(int nameId, ref NativeList<int> data)
        {
            AppendInt(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void AppendInt(int nameId, int* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Int);
            }
        }

        public void RemoveRangeInt(int nameId, int startIndex, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.RemoveRange(startIndex * GPUDefine.Stride_Int, count * GPUDefine.Stride_Int);
            }
        }

        #endregion
        
        #region Float

        public void SetFloat(int nameId, float[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float* ptr = &data[0])
                {
                    property.nativeList.Clear();
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float);
                }
            }
        }

        public void SetFloat(int nameId, List<float> data)
        {
            SetFloat(nameId, data.ToArray());
        }
        
        public void SetFloat(int nameId, ref NativeArray<float> data)
        {
            SetFloat(nameId, (float*) data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat(int nameId, ref NativeList<float> data)
        {
            SetFloat(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat(int nameId, float* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.Clear();
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float);
            }
        }

        public void AppendFloat(int nameId, float[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float* ptr = &data[0])
                {
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float);
                }
            }
        }

        public void AppendFloat(int nameId, List<float> data)
        {
            AppendFloat(nameId, data.ToArray());
        }

        public void AppendFloat(int nameId, ref NativeArray<float> data)
        {
            AppendFloat(nameId, (float*) data.GetUnsafePtr(), data.Length);
        }
        
        public void AppendFloat(int nameId, ref NativeList<float> data)
        {
            AppendFloat(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void AppendFloat(int nameId, float* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float);
            }
        }

        public void RemoveRangeFloat(int nameId, int startIndex, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.RemoveRange(startIndex * GPUDefine.Stride_Float, count * GPUDefine.Stride_Float);
            }
        }

        #endregion

        #region Float2

        public void SetFloat2(int nameId, float2[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float2* ptr = &data[0])
                {
                    property.nativeList.Clear();
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float2);
                }
            }
        }

        public void SetFloat2(int nameId, List<float2> data)
        {
            SetFloat2(nameId, data.ToArray());
        }
        
        public void SetFloat2(int nameId, ref NativeArray<float2> data)
        {
            SetFloat2(nameId, (float2*) data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat2(int nameId, ref NativeList<float2> data)
        {
            SetFloat2(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat2(int nameId, float2* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.Clear();
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float2);
            }
        }

        public void AppendFloat2(int nameId, float2[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float2* ptr = &data[0])
                {
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float2);
                }
            }
        }

        public void AppendFloat2(int nameId, List<float2> data)
        {
            AppendFloat2(nameId, data.ToArray());
        }

        public void AppendFloat2(int nameId, ref NativeArray<float2> data)
        {
            AppendFloat2(nameId, (float2*) data.GetUnsafePtr(), data.Length);
        }
        
        public void AppendFloat2(int nameId, ref NativeList<float2> data)
        {
            AppendFloat2(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void AppendFloat2(int nameId, float2* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float2);
            }
        }

        public void RemoveRangeFloat2(int nameId, int startIndex, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.RemoveRange(startIndex * GPUDefine.Stride_Float2, count * GPUDefine.Stride_Float2);
            }
        }

        #endregion

        #region Float3

        public void SetFloat3(int nameId, float3[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float3* ptr = &data[0])
                {
                    property.nativeList.Clear();
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float3);
                }
            }
        }

        public void SetFloat3(int nameId, List<float3> data)
        {
            SetFloat3(nameId, data.ToArray());
        }
        
        public void SetFloat3(int nameId, ref NativeArray<float3> data)
        {
            SetFloat3(nameId, (float3*) data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat3(int nameId, ref NativeList<float3> data)
        {
            SetFloat3(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat3(int nameId, float3* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.Clear();
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float3);
            }
        }

        public void AppendFloat3(int nameId, float3[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float3* ptr = &data[0])
                {
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float3);
                }
            }
        }

        public void AppendFloat3(int nameId, List<float3> data)
        {
            AppendFloat3(nameId, data.ToArray());
        }

        public void AppendFloat3(int nameId, ref NativeArray<float3> data)
        {
            AppendFloat3(nameId, (float3*) data.GetUnsafePtr(), data.Length);
        }
        
        public void AppendFloat3(int nameId, ref NativeList<float3> data)
        {
            AppendFloat3(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void AppendFloat3(int nameId, float3* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float3);
            }
        }

        public void RemoveRangeFloat3(int nameId, int startIndex, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.RemoveRange(startIndex * GPUDefine.Stride_Float3, count * GPUDefine.Stride_Float3);
            }
        }

        #endregion
        
        #region Float4

        public void SetFloat4(int nameId, float4[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float4* ptr = &data[0])
                {
                    property.nativeList.Clear();
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float4);
                }
            }
        }

        public void SetFloat4(int nameId, List<float4> data)
        {
            SetFloat4(nameId, data.ToArray());
        }
        
        public void SetFloat4(int nameId, ref NativeArray<float4> data)
        {
            SetFloat4(nameId, (float4*) data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat4(int nameId, ref NativeList<float4> data)
        {
            SetFloat4(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat4(int nameId, float4* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.Clear();
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float4);
            }
        }

        public void AppendFloat4(int nameId, float4[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float4* ptr = &data[0])
                {
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float4);
                }
            }
        }

        public void AppendFloat4(int nameId, List<float4> data)
        {
            AppendFloat4(nameId, data.ToArray());
        }

        public void AppendFloat4(int nameId, ref NativeArray<float4> data)
        {
            AppendFloat4(nameId, (float4*) data.GetUnsafePtr(), data.Length);
        }
        
        public void AppendFloat4(int nameId, ref NativeList<float4> data)
        {
            AppendFloat4(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void AppendFloat4(int nameId, float4* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float4);
            }
        }

        public void RemoveRangeFloat4(int nameId, int startIndex, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.RemoveRange(startIndex * GPUDefine.Stride_Float4, count * GPUDefine.Stride_Float4);
            }
        }

        #endregion
        
        #region Float4x4

        public void SetFloat4x4(int nameId, float4x4[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float4x4* ptr = &data[0])
                {
                    property.nativeList.Clear();
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float4x4);
                }
            }
        }

        public void SetFloat4x4(int nameId, List<float4x4> data)
        {
            SetFloat4x4(nameId, data.ToArray());
        }
        
        public void SetFloat4x4(int nameId, ref NativeArray<float4x4> data)
        {
            SetFloat4x4(nameId, (float4x4*) data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat4x4(int nameId, ref NativeList<float4x4> data)
        {
            SetFloat4x4(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void SetFloat4x4(int nameId, float4x4* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.Clear();
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float4x4);
            }
        }

        public void AppendFloat4x4(int nameId, float4x4[] data)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                fixed (float4x4* ptr = &data[0])
                {
                    property.nativeList.AddRange(ptr, data.Length * GPUDefine.Stride_Float4x4);
                }
            }
        }

        public void AppendFloat4x4(int nameId, List<float4x4> data)
        {
            AppendFloat4x4(nameId, data.ToArray());
        }

        public void AppendFloat4x4(int nameId, ref NativeArray<float4x4> data)
        {
            AppendFloat4x4(nameId, (float4x4*) data.GetUnsafePtr(), data.Length);
        }
        
        public void AppendFloat4x4(int nameId, ref NativeList<float4x4> data)
        {
            AppendFloat4x4(nameId, data.GetUnsafePtr(), data.Length);
        }

        public void AppendFloat4x4(int nameId, float4x4* ptr, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.AddRange(ptr, count * GPUDefine.Stride_Float4x4);
            }
        }

        public void RemoveRangeFloat4x4(int nameId, int startIndex, int count)
        {
            var property = GetPropertyData(nameId);
            if (property != null)
            {
                property.nativeList.RemoveRange(startIndex * GPUDefine.Stride_Float4x4, count * GPUDefine.Stride_Float4x4);
            }
        }

        #endregion
        
        #endregion
        
        public void Draw()
        {
            if (_propertyDataList.Count == 0)
            {
                return;
            }
            var count = _propertyDataList[0].nativeList.Length / _propertyDataList[0].stride;
            if (count == 0)
            {
                return;
            }
            
            var batchCount = (int) math.ceil((float) count / (float) _maximum);
            for (var i = 0; i < batchCount; i++)
            {
                var startIndex = i * _maximum;
                var len = _maximum;
                if (i == batchCount - 1)
                {
                    len = count - (_maximum * i);
                }
                
                // 缺少的部分需要创建
                if (_batchDataList.Count <= i)
                {
                    _batchDataList.Add(Recyclable.Get<GPUBatchData>());
                }
                
                // 拷贝数据
                var batchData = _batchDataList[i];
                for (var j = 0; j < _propertyDataList.Count; j++)
                {
                    var propertyData = _propertyDataList[j];
                    var nameId = propertyData.nameId;
                    var batchPropertyData = batchData.GetOrCreateData(nameId, _maximum, propertyData.stride);
                    var destPtr = batchPropertyData.nativeArray.GetUnsafePtr();
                    var sourcePtr = propertyData.nativeList.GetUnsafePtr();
                    var source = sourcePtr + startIndex * propertyData.stride;
                    UnsafeUtility.MemCpy(destPtr, source, batchPropertyData.length);
                    batchData.block.SetBuffer(nameId, batchData.GetBuffer(nameId));
                }
                // 设置参数
                batchData.args[0] = _meshIndexCount; // 每个实例的索引数
                batchData.args[1] = (uint) len; // 实例数
                batchData.args[2] = _meshIndexStart; // 起始索引位置
                batchData.args[3] = _meshBaseVertex; // 基顶点位置
                batchData.args[4] = 0; // 起始实例位置
                
                // 提交绘制
                Graphics.DrawMeshInstancedIndirect(_mesh, 0, _material, BOUNDS, batchData.argsBuf, 0, batchData.block, _castShadows, _receiveShadows, _layer, _camera);
            }

            // 超出部分回收
            if (_batchDataList.Count > batchCount)
            {
                for (var i = batchCount; i < _batchDataList.Count; i++)
                {
                    Recyclable.Recycle(_batchDataList[i]);
                }
                _batchDataList.RemoveRange(batchCount, _batchDataList.Count - batchCount);
            }
        }

        public void Dispose()
        {
            for (var i = 0; i < _propertyDataList.Count; i++)
            {
                var propertyData = _propertyDataList[i];
                propertyData.nativeList.Dispose();
            }
            _propertyDataList.Clear();
            for (var i = 0; i < _batchDataList.Count; i++)
            {
                Recyclable.Recycle(_batchDataList[i]);
            }
            _batchDataList.Clear();
        }
    }

    /// <summary>
    /// 设置数据，长度可变
    /// </summary>
    public class GPUPropertyData
    {
        public int nameId;
        public int stride;
        public NativeList<byte> nativeList;
    }

    /// <summary>
    /// 批次提交数据
    /// </summary>
    public class GPUBatchData : ICacheable
    {
        private MaterialPropertyBlock _block;
        private uint[] _args;
        private GraphicsBuffer _argsBuf;
        
        private List<GPUBatchPropertyData> _dataList = new List<GPUBatchPropertyData>();
        
        public GPUBatchData()
        {
            _block = new MaterialPropertyBlock();
            _args = new uint[5] {0, 0, 0, 0, 0};
            _argsBuf = new GraphicsBuffer(GraphicsBuffer.Target.IndirectArguments, 1, 5 * sizeof(uint));
        }

        public MaterialPropertyBlock block => _block;

        public uint[] args => _args;

        public GraphicsBuffer argsBuf
        {
            get
            {
                _argsBuf.SetData(_args);
                return _argsBuf;
            }
        }

        private GPUBatchPropertyData GetPropertyData(int nameId)
        {
            for (var i = 0; i < _dataList.Count; i++)
            {
                if (_dataList[i].nameId == nameId)
                {
                    return _dataList[i];
                }
            }
            return null;
        }

        public GPUBatchPropertyData GetOrCreateData(int nameId, int count, int stride)
        {
            var data = GetPropertyData(nameId);
            if (data != null)
            {
                return data;
            }

            var length = count * stride;
            var buffer = new GraphicsBuffer(GraphicsBuffer.Target.Structured, count, stride);
            
            data = new GPUBatchPropertyData();
            data.nameId = nameId;
            data.length = length;
            data.nativeArray = new NativeArray<byte>(length, Allocator.Persistent);
            data.buffer = buffer;
            _dataList.Add(data);
            return data;
        }

        public GraphicsBuffer GetBuffer(int nameId)
        {
            var data = GetPropertyData(nameId);
            data.buffer.SetData(data.nativeArray);
            return data.buffer;
        }
        
        public void OnRecycle()
        {
            _block.Clear();
            for (var i = 0; i < 5; i++)
            {
                _args[i] = 0;
            }
            for (var i = 0; i < _dataList.Count; i++)
            {
                var data = _dataList[i];
                data.nativeArray.Dispose();
                data.buffer.Dispose();
            }
            _dataList.Clear();
        }

        public void OnReuse()
        {
        }
    }

    /// <summary>
    /// 提交数据，长度固定
    /// </summary>
    public class GPUBatchPropertyData
    {
        public int nameId;
        public int length;
        public NativeArray<byte> nativeArray;
        public GraphicsBuffer buffer;
    }
}
